#include "ulib.h"
#include "sysenv.h"
#include "syscall.h"

struct va_item_t
{
	 unsigned int flag : 2;
	 unsigned int idx  : 30;
};

#define VA_FLAG_FH   0
#define VA_FLAG_FB   1
#define VA_FLAG_FT   2
#define VA_FLAG_USED 3

struct va_ctrl_t
{
	 spin_lock_t lock;
	 uint32_t    free_head;
};

static char __va_item_size_assert[sizeof(struct va_item_t) == 4 ? 0 : -1];

extern volatile struct va_ctrl_t va_ctrl;
extern volatile struct va_item_t va_items[PAGETAB_COUNT * PAGEDIR_COUNT];

static spin_lock_t mm_lock;

static int
page_alloc(uintptr_t vaddr)
{
	 physaddr_t p;
	 int result;
	 if ((result = sc_ppage_alloc(1, 0, &p)) < 0)
	 {
		  return result;
	 }
	 if ((result = sc_mmap_set(cur_pid(), (void *)vaddr, p, SC_MMAP_SET_PERM_WRITE)) < 0)
	 {
		  sc_ppage_free(p);
	 }
	 return result;
}

static int
page_free(uintptr_t vaddr)
{
	 sc_ppage_free(vpt[PAGE_NUM(vaddr)]);
	 sc_mmap_set(cur_pid(), (void *)vaddr, 0, 0);
}

static int
va_init(void)
{
	 int result;
	 int i;
	 
	 for (i = 0; i != (1 << 10); ++ i)
	 {
		  if ((result = page_alloc(UHHEAD + (i << PAGE_SHIFT))) < 0)
			   return result;
	 }

	 spl_init(&va_ctrl.lock);
	 va_ctrl.free_head = PAGE_NUM(UHEAP);

	 result = 0;

	 int head = PAGE_NUM(UHEAP);
	 int tail = PAGE_NUM(UWLIM) - 1;
	 
	 va_items[head].flag = VA_FLAG_FH;
	 va_items[head].idx	 = 0;
	 va_items[tail].flag = VA_FLAG_FT;
	 va_items[tail].idx  = 0;

	 va_items[head + 1].flag = VA_FLAG_FB;
	 va_items[head + 1].idx  = tail - 1;
	 va_items[tail - 1].flag = VA_FLAG_FB;
	 va_items[tail - 1].idx  = head + 1;

	 return result;
}

static void *
va_alloc(uint32_t size)
{
	 if (size & (PAGE_SIZE - 1))
		  size = (size >> PAGE_SHIFT) + 1;
	 else size = (size >> PAGE_SHIFT);
	 /* minimum size of 2 * PAGE_SIZE */
	 if (size < 2) size = 2;

	 // sc_kprintf("va_alloc page size = %d\n", size);
	 
	 int result, last_tail, tail;
	 spl_acquire(&va_ctrl.lock);
	 last_tail = 0;
	 result = va_ctrl.free_head;
	 while (result != 0)
	 {
		  /* Get tail of free block */
		  tail = va_items[result + 1].flag == VA_FLAG_FB ?
			   va_items[result + 1].idx + 1 : result + 1;

		  // sc_kprintf("Looking at %d - %d\n", result, tail);
		  
		  uint32_t _size = tail - result + 1;
		  /* Do not alloc the 1 page size fragment */
		  if (_size == size || _size == size + 1)
		  {
			   /* Remove the free block and catch the result */
			   if (last_tail == 0)
					va_ctrl.free_head = va_items[tail].idx;
			   else va_items[last_tail].idx = va_items[tail].idx;

			   if (va_items[tail].idx != 0)
					va_items[va_items[tail].idx].idx = 0;

			   break;
		  }
		  else if (_size > size)
		  {
			   /* Split the free block then break */
			   /* The new tail */
			   int n_tail = tail - size;
			   if (n_tail - 1 != result)
			   {
					va_items[n_tail - 1].flag = VA_FLAG_FB;
					va_items[n_tail - 1].idx  = result + 1;
			   }
			   va_items[n_tail] = va_items[tail];
			   if (result + 1 < n_tail - 1)
					va_items[result + 1].idx = n_tail - 1;
			   result = n_tail + 1;

			   break;
		  }
		  else
		  {
			   /* Not enough space, try next one on the list */
			   result = va_items[tail].idx;
			   last_tail = tail;
		  }
	 }
	 spl_release(&va_ctrl.lock);

	 // sc_kprintf("va_alloc ends with result idx = %d\n", result);
	 if (result != 0)
	 {
		  va_items[result].flag = VA_FLAG_USED;
		  va_items[result].idx  = tail;
		  va_items[tail].flag   = VA_FLAG_USED;
		  va_items[tail].idx    = result;

		  // sc_kprintf("result = %d tail = %d\n", result, tail);

		  return (void *)(result << PAGE_SHIFT);
	 }
	 else
	 {
		  sc_kprintf("va_alloc failed for size %d!\n", size);
		  return NULL;
	 }
}

uint32_t
va_psize(void *addr)
{
	 int head = PAGE_NUM(addr);
	 int tail = va_items[head].idx;
	 
	 return tail - head;
}



static void
va_free(void *addr)
{
	 int head = PAGE_NUM(addr);
	 int tail = va_items[head].idx;

	 // sc_kprintf("va_free %d - %d\n", head, tail);

	 spl_acquire(&va_ctrl.lock);
	 /* Merge up */
	 if (head > PAGE_NUM(UHEAP) &&
		 va_items[head - 1].flag != VA_FLAG_USED)
	 {
		  va_items[tail] = va_items[head - 1];
		  head = va_items[head - 2].flag == VA_FLAG_FB ?
			   va_items[head - 2].idx - 1 : head - 2;

		  // sc_kprintf("merge up! head = %d,\n", head);
	 }
	 else
	 {
		  va_items[head].flag = VA_FLAG_FH;
		  va_items[head].idx  = 0;
		  
		  va_items[tail].flag = VA_FLAG_FT;
		  va_items[tail].idx  = 0;
	 }

	 /* Merge down */
	 if (tail + 1 < PAGE_NUM(UWLIM) &&
		 va_items[tail + 1].flag != VA_FLAG_USED)
	 {
		  int n_tail = va_items[tail + 2].flag == VA_FLAG_FB ?
			   va_items[tail + 2].idx + 1: tail + 2;

		  /* remove block from list */
		  if (va_items[tail + 1].idx == 0)
			   va_ctrl.free_head = va_items[n_tail].idx;
		  else va_items[va_items[tail + 1].idx].idx = va_items[n_tail].idx;
		  if (va_items[n_tail].idx != 0)
			   va_items[va_items[n_tail].idx].idx = va_items[tail + 1].idx;

		  va_items[n_tail] = va_items[tail];
		  tail = n_tail;

		  // sc_kprintf("merge down! tail = %d,\n", tail);
	 }

	 if (head == PAGE_NUM(addr))
	 {
		  /* insert into list */
		  va_items[head].idx = 0;
		  if ((va_items[tail].idx = va_ctrl.free_head) != 0)
			   va_items[va_ctrl.free_head].idx = head;
		  va_ctrl.free_head = head;
	 }

	 /* Set the inside block info */
	 if (head + 1 != tail)
	 {
		  va_items[head + 1].flag = VA_FLAG_FB;
		  va_items[head + 1].idx  = tail - 1;
		  va_items[tail - 1].flag = VA_FLAG_FB;
		  va_items[tail - 1].idx  = head + 1;
	 }
	 spl_release(&va_ctrl.lock);

	 // sc_kprintf("va_free finished\n");
}

/* ============================================================ */

void *
bmalloc(uint32_t size)
{
	 char *result = va_alloc(size);
	 if (result != NULL)
	 {
		  char *cur;
		  for (cur = result; cur < result + size; cur += PAGE_SIZE)
		  {
			   page_alloc((uintptr_t)cur);
		  }
	 }


	 return result;
}

void
bfree(void *addr)
{
	 if (addr == NULL) return;
	 
	 int psize = va_psize(addr);
	 
	 char *cur = addr, *end = cur + (psize << PAGE_SHIFT);
	 for ( ; cur != end; cur += PAGE_SIZE)
	 {
		  page_free((uintptr_t)cur);
	 }

	 va_free(addr);
}

/* ============================================================ */

#define EKL_META_SIZE 4
#define EKL_MIN_ALLOC 8
#define EKL_MIN_SHIFT 3
#define EKL_ALLOC_DELTA_SHIFT 10
#define EKL_MAX_SHIFT (EKL_MIN_SHIFT + EKL_ALLOC_DELTA_SHIFT)
#define EKL_MAX_ALLOC (EKL_MIN_ALLOC << EKL_ALLOC_DELTA_SHIFT)
#define EKL_MIN_BLOCK_PAGE 4
#define EKL_MAX_BLOCK_PAGE 256

static char __ekl_size_assert[(EKL_MIN_BLOCK_PAGE << PAGE_SHIFT) >= EKL_MAX_ALLOC ? 0 : -1];

struct ekl_ctrl_t
{
	 uint16_t alloc_size;
	 uint16_t block_size;
	 uint32_t head;
};

struct ekl_ctrl_t ekl_ctrls[EKL_ALLOC_DELTA_SHIFT + 1];

int
ekl_init(void)
{
	 int i;
	 for (i = 0; i <= EKL_ALLOC_DELTA_SHIFT; ++ i)
	 {
		  ekl_ctrls[i].alloc_size = EKL_MIN_ALLOC << i;
		  ekl_ctrls[i].head = -1;
		  ekl_ctrls[i].block_size = EKL_MIN_BLOCK_PAGE;
	 }
}

void *
ekl_ialloc(struct ekl_ctrl_t *ctrl)
{
	 if (ctrl->head == -1)
	 {
		  uint32_t size = ctrl->block_size << PAGE_SHIFT;
		  if (ctrl->block_size < EKL_MAX_BLOCK_PAGE)
			   ctrl->block_size <<= 1;

		  void *block = bmalloc(size);

		  if (block == NULL)
		  {
			   return NULL;
		  }
		  
		  memset(block, 0, size);

		  uint32_t *tail = (uint32_t *)(
			   (uintptr_t)block +
			   ((size / ctrl->alloc_size - 1) * (ctrl->alloc_size)));
		  *tail = -1;

		  ctrl->head = (uintptr_t)block;
	 }
	 
	 uint32_t *head = (uint32_t *)ctrl->head;
	 ctrl->head = (*head == 0) ? (ctrl->head + ctrl->alloc_size) : (*head);

	 *head = (uintptr_t)ctrl;

	 return head + 1;
}

void *
ekl_alloc(uint32_t size)
{
	 spl_acquire(&mm_lock);	 
	 int size_index = bsr(size + EKL_META_SIZE - 1) - EKL_MIN_SHIFT + 1;
	 
	 void *result;
	 if (size_index < 0 || size_index > EKL_ALLOC_DELTA_SHIFT)
		  result = NULL;
	 else result = ekl_ialloc(ekl_ctrls + size_index);
	 
	 spl_release(&mm_lock);	 
	 return result;
}

void
ekl_free(void *ptr)
{
	 if (ptr == NULL) return;
	 spl_acquire(&mm_lock);
	 
	 uint32_t *head = ((uint32_t *)ptr - 1);
	 struct ekl_ctrl_t *ctrl = (struct ekl_ctrl_t *)*head;

	 *head = ctrl->head;
	 ctrl->head = (uintptr_t)head;

	 spl_release(&mm_lock);
}

/* ============================================================ */

int
mm_init(void)
{
	 int result;
	 
	 if ((result = va_init()) < 0) return result;
	 if ((result = ekl_init()) < 0) return result;

	 spl_init(&mm_lock);
}

void
*malloc(uint32_t size)
{
	 if (size == 0) return NULL;
	 if (size < 4) size = 4;

	 void *result;
	 if (size + EKL_META_SIZE > EKL_MAX_ALLOC)
		  result = bmalloc(size);
	 else result = ekl_alloc(size);

	 return result;
}

void
free(void *ptr)
{
	 if (ptr == NULL) return;
	 if (((uintptr_t)ptr & (PAGE_SIZE - 1)) == 0)
	 {
		  bfree(ptr);
	 }
	 else ekl_free(ptr);
}

void *
pmalloc(uint32_t size, int dma)
{
	 size = ((size - 1) | (PAGE_SIZE - 1)) + 1;
	 
	 char *result = va_alloc(size);
	 if (result != NULL)
	 {
		  physaddr_t pstart;
		  if (sc_ppage_alloc(size >> PAGE_SHIFT, dma, &pstart) != 0)
		  {
			   va_free(result);
			   result = NULL;
		  }
		  else
		  {
			   char *cur = result;
			   physaddr_t pcur = pstart;
			   for (; cur < result + size; cur += PAGE_SIZE)
			   {
					if (sc_mmap_set(cur_pid(), (void *)cur, pcur, SC_MMAP_SET_PERM_WRITE) < 0)
					{
						 sc_ppage_free(pstart);
						 va_free(result);
						 result = NULL;
						 break;
					}
					pcur += PAGE_SIZE;
			   }
		  }
	 }
	 
	 return result;
}

void
pfree(void *addr)
{
	 if (addr == NULL) return;
	 int psize = va_psize(addr);
	 physaddr_t pstart = PTE_ADDR(vpt[PAGE_NUM(addr)]);
		  
	 char *cur = addr, *end = cur + (psize << PAGE_SHIFT);
	 for ( ; cur != end; cur += PAGE_SIZE)
	 {
		  sc_mmap_set(cur_pid(), (void *)cur, 0, 0);
	 }
	 va_free(addr);
	 sc_ppage_free(pstart);
}

void *
vmalloc(uint32_t size, physaddr_t start)
{
	 size = ((size - 1) | (PAGE_SIZE - 1)) + 1;
	 
	 char *result = va_alloc(size);
	 if (result != NULL)
	 {
		  char *cur = result;
		  physaddr_t pcur = start;
		  for (; cur < result + size; cur += PAGE_SIZE)
		  {
			   if (sc_mmap_set(cur_pid(), (void *)cur, pcur, SC_MMAP_SET_PERM_WRITE) < 0)
			   {
					va_free(result);
					result = NULL;
					break;
			   }
			   pcur += PAGE_SIZE;
		  }
	 }
	 
	 return result;
}

void
vfree(void *addr)
{
	 if (addr == NULL) return;
	 int psize = va_psize(addr);
		  
	 char *cur = addr, *end = cur + (psize << PAGE_SHIFT);
	 for ( ; cur != end; cur += PAGE_SIZE)
	 {
		  sc_mmap_set(cur_pid(), (void *)cur, 0, 0);
	 }

	 va_free(addr);
}

void
self_mmap_set(uintptr_t vaddr, physaddr_t paddr, uint32_t flags)
{
	 sc_mmap_set(cur_pid(), (void *)vaddr, paddr, flags ? SC_MMAP_SET_PERM_WRITE : 0);
}

physaddr_t
self_mmap_get(uintptr_t vaddr, uint32_t *flags)
{
	 uint32_t pte;
	 if (!(vpd[PAGEDIR_IDX(vaddr)] & PTE_P))
		  pte = 0;
	 else pte = vpt[PAGE_NUM(vaddr)];

	 if (flags != NULL) *flags = (pte & PTE_W) ? MMAP_W : 0;
	 return PTE_ADDR(pte);
}

static void *proc_datas[PROCESS_MAX_COUNT];

void
proc_set_data(void *data)
{
	 proc_datas[cur_pid()] = data;
}
	 
void *
proc_get_data(void)
{
	 return proc_datas[cur_pid()];
}
